package com.lu.component.http.proxy.support.flow;

import cn.hutool.core.lang.Assert;
import com.lu.component.rpc.annotations.RpcApiRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

/**
 * 单机流控规则
 *
 * @author: zhanglu
 * @date: 2023/8/31 17:34
 * @modify: 修改人信息, 修改内容 2023/8/31 17:34
 */
public class RpcFlowControlStandalone extends RpcFlowControlAbstract{

    private static final Logger log = LoggerFactory.getLogger(RpcFlowControlStandalone.class);

    public static final Map<String, Semaphore> SEMAPHORE = new ConcurrentHashMap<>();

    public RpcFlowControlStandalone(RpcApiRequest rpcApiRequest) {
        super(rpcApiRequest);
    }

    @Override
    public <T> T execute(String key, Supplier<T> supplier, Supplier<T> fallback) {
        Semaphore semaphore = getSemaphore(key, rpcApiRequest);
        try {
            Assert.isTrue(semaphore.tryAcquire(rpcApiRequest.limitWaitTime(), TimeUnit.MILLISECONDS), () -> new InterruptedException("服务限流，令牌获取等待超时"));
            log.debug("【单机限流】获取到令牌，key：{}", key);
            T t = supplier.get();
            log.debug("【单机限流】rpc请求执行完成，key：{}", key);
            semaphore.release();
            return t;
        } catch (InterruptedException e) {
            log.debug(e.getMessage());
            return fallback.get();
        } catch (Exception e) {
            log.error("【单机限流】rpc请求异常", e);
            return fallback.get();
        }
    }

    /**
     * 原子性操作 - 更新流控规则
     * @param key
     */
    protected static synchronized Semaphore getSemaphore(String key, RpcApiRequest rpcApiRequest){
        Semaphore semaphore = SEMAPHORE.get(key);
        if(semaphore != null){
            return semaphore;
        }else {
            synchronized (RpcFlowControlStandalone.class){
                if(SEMAPHORE.containsKey(key)){
                    return SEMAPHORE.get(key);
                } else {
                    log.info("【单机限流】加载流控规则，key：{}，配额：{}", key, rpcApiRequest.permits());
                    SEMAPHORE.put(key, new Semaphore(rpcApiRequest.permits()));
                    return SEMAPHORE.get(key);
                }
            }
        }
    }

}
